package org.raftkv.raft.core;

import org.raftkv.raft.core.storage.impl.MemoryStorage;
import org.raftkv.raft.exceptions.RaftException;
import org.raftkv.raft.pb.Message;
import org.raftkv.raft.pb.enums.MessageType;
import org.raftkv.raft.tracker.Progress;
import org.raftkv.raft.tracker.ProgressTracker;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
import java.util.function.Consumer;
import java.util.function.Function;

import static org.raftkv.raft.core.TestUtils.*;

public class Network {

    Map<Long, StateMachine> peers;

    Map<Long, MemoryStorage> storage;

    Map<Connem, Double> dropm;

    HashMap<MessageType, Boolean> ignorem;

    Function<Message, Boolean> msgHook;

    public static Network newNetwork(StateMachine... peers) {
        return newNetworkWithConfig(null, peers);
    }

    public static Network newNetworkWithConfig(Consumer<Config> configFunc, StateMachine... peers) {
        if (peers == null) {
            peers = new StateMachine[1];
            peers[0] = null;
        }
        int size = peers.length;
        List<Long> peerAddrs = TestUtils.idsBySize(size);

        Map<Long, StateMachine> npeers = new HashMap<>();
        Map<Long, MemoryStorage> nstorage = new HashMap<>();

        for (int j = 0; j < size; j++) {
            StateMachine p = peers[j];
            Long id = peerAddrs.get(j);

            if (p instanceof Raft) {
                Raft v = (Raft) p;
                Map<Long, Boolean> learners = new HashMap<>();
                Optional.ofNullable(v.getPrs().getConfig().getLearners())
                        .ifPresent((e) -> {
                            for (Long i : e) {
                                learners.put(i, true);
                            }
                        });
                v.setId(id);
                v.setPrs(ProgressTracker.makeProgressTracker(v.getPrs().getMaxInflight()));
                if (learners.size() > 0) {
                    v.getPrs().getConfig().setLearners(new HashSet<>());
                }
                for (int i = 0; i < size; i++) {
                    Progress pr = new Progress();
                    boolean ok = learners.containsKey(peerAddrs.get(i));
                    if (ok) {
                        pr.setLearner(true);
                        v.getPrs().getConfig().getLearners().add(peerAddrs.get(i));
                    } else {
                        v.getPrs().getConfig().getVoters().CONFIGS[0].add(peerAddrs.get(i));
                    }
                    v.getPrs().getProgress().getProgressMap().put(peerAddrs.get(i), pr);
                }
                v.reset(v.getTerm());
                npeers.put(id, v);
            } else if (p instanceof BlackHole) {
                BlackHole v = (BlackHole) p;
                npeers.put(id, v);
            } else if (p == null) {
                nstorage.put(id, newTestMemoryStorage(withPeers(peerAddrs)));
                Config cfg = newTestConfig(id, 10, 1, nstorage.get(id));
                if (configFunc != null) {
                    configFunc.accept(cfg);
                }
                Raft sm = Raft.newRaft(cfg);
                npeers.put(id, sm);
            } else {
                throw new RaftException(String.format("unexpected state machine type: %s", p.getClass().getTypeName()));
            }
        }
        Network n = new Network();
        n.peers = npeers;
        n.storage = nstorage;
        n.dropm = new HashMap<>();
        n.ignorem = new HashMap<>();
        return n;
    }

    public void send(List<Message> ms) {
        List<Message> msgs = new ArrayList<>(ms);
        while (msgs.size() > 0) {
            Message m = msgs.get(0);
            StateMachine p = this.peers.get(m.getTo());
            p.step(m);
            List<Message> filteredMsgs = this.filter(p.readMessages());
            msgs = msgs.subList(1, msgs.size());
            msgs.addAll(filteredMsgs);
        }
    }

    /**
     * 模拟在网络中发送消息
     */
    public void send(Message... ms) {
        this.send(Arrays.stream(ms).toList());
    }

    /**
     * 模拟丢包，丢包率由 prec 参数指定
     */
    public void drop(long from, long to, double prec) {
        this.dropm.put(new Connem(from, to), prec);
    }

    /**
     * 模拟在网络上切断连接
     */
    public void cut(long one, long other) {
        drop(one, other, 2.0);
        drop(other, one, 2.0);
    }

    /**
     * 模拟在网络上隔离节点
     */
    public void isolate(long id) {
        if (peers != null) {
            for (int i = 0; i < peers.size(); i++) {
                long nid = i + 1;
                if (nid != id) {
                    drop(id, nid, 1.0);
                    drop(nid, id, 1.0);
                }
            }
        }
    }

    /**
     * 模拟在网络上忽略特定类型的消息
     */
    public void ignore(MessageType t) {
        this.ignorem.put(t, true);
    }

    /**
     * 重置网络环境
     */
    public void recover() {
        this.dropm = new HashMap<>();
        this.ignorem = new HashMap<>();
    }

    /**
     * 消息过滤
     */
    public List<Message> filter(List<Message> msgs) {
        List<Message> mm = new ArrayList<>();

        if (msgs != null) {
            for (Message m : msgs) {
                if (this.ignorem.get(m.getType()) != null && this.ignorem.get(m.getType())) {
                    continue;
                }
                switch (m.getType()) {
                    case MsgHup:
                        throw new RaftException("unexpected msgHup");
                    default:
                        Double perc = this.dropm.get(new Connem(m.getFrom(), m.getTo()));
                        double n = new Random().nextDouble();
                        if (perc != null && n < perc) {
                            continue;
                        }
                }
                if (this.msgHook != null) {
                    if (!this.msgHook.apply(m)) {
                        continue;
                    }
                }
                mm.add(m);
            }
        }

        return mm;
    }

}
